# -*- coding: utf-8 -*-

# Copyright (c) Huoty - All rights reserved
# Author: Huoty <sudohuoty@gmail.com>
# CreateTime: 2016-11-25 11:52:59

"""一些实用的函数和类的封装"""

from __future__ import absolute_import

import re
import ast
import types
import inspect
import signal
import functools
import contextlib
from collections import OrderedDict


class TimeOutError(BaseException):
    """Raised when a timeout happens"""

    def __init__(self, msg='TimeOutError'):
        BaseException.__init__(self, msg)


def timeout(timeout):
    """Return a decorator that raises a TimeOutError exception

    after timeout seconds, if the decorated function did not return.
    """
    def dec(f):
        @functools.wraps(f)
        def wrapper(*args, **kwargs):
            def handler(signum, frame):
                raise TimeOutError()

            old_handler = signal.signal(signal.SIGALRM, handler)
            signal.alarm(timeout)

            try:
                return f(*args, **kwargs)  # f() always returns, in this scheme
            finally:
                signal.signal(signal.SIGALRM, old_handler)  # Old signal handler is restored
                signal.alarm(0)  # Alarm removed
        return wrapper
    return dec


def timed_run(seconds, func, *args, **kwargs):
    return timeout(seconds)(func)(*args, **kwargs)


class context_timeout(object):

    def __init__(self, seconds=3, error_message='Timeout'):
        self.seconds = seconds
        self.error_message = error_message

        self._old_handler = None

    def handle_timeout(self, signum, frame):
        raise TimeOutError(self.error_message)

    def __enter__(self):
        self._old_handler = signal.signal(signal.SIGALRM, self.handle_timeout)
        signal.alarm(self.seconds)

    def __exit__(self, type, value, traceback):
        signal.alarm(0)
        signal.signal(signal.SIGALRM, self._old_handler)


def _make_assert_message(frame, regex):
    def extract_condition():
        code_context = inspect.getframeinfo(frame)[3]
        if not code_context:
            return ''
        match = re.search(regex, code_context[0])
        if not match:
            return ''
        return match.group(1).strip()

    class ReferenceFinder(ast.NodeVisitor):
        def __init__(self):
            self.names = []
        def find(self, tree, frame):
            self.visit(tree)
            nothing = object()
            deref = OrderedDict()
            for name in self.names:
                value = frame.f_locals.get(name, nothing) or frame.f_globals.get(name, nothing)
                if value is not nothing and not isinstance(value, (types.ModuleType, types.FunctionType)):
                    deref[name] = repr(value)
            return deref
        def visit_Name(self, node):
            self.names.append(node.id)

    condition = extract_condition()
    if not condition:
        return
    deref = ReferenceFinder().find(ast.parse(condition), frame)
    deref_str = ''
    if deref:
        deref_str = ' with ' + ', '.join('{}={}'.format(k, v) for k, v in deref.items())
    return 'assertion {} failed{}'.format(condition, deref_str)


def affirm(condition, message=None):
    if condition:
        return

    if message:
        raise AssertionError(str(message))

    frame = inspect.currentframe().f_back
    regex = r'affirm\s*\(\s*(.+)\s*\)'
    message = _make_assert_message(frame, regex)

    raise AssertionError(message)


@contextlib.contextmanager
def ignored(*exceptions):
    """A context managers ignored some exceptions

    Usage:
    ------
        with ignored(OSError):
            os.remove('i_probably_do_not_exist')

    Reference http://www.rosipov.com/blog/python-ignored-context-manager/
    """
    try:
        yield
    except exceptions:
        pass


class LazyEval(object):

    def __init__(self, func, *args, **kwargs):
        self.func = func
        self.args = args
        self.kwargs = kwargs

    def __str__(self):
        return str(self.func(*self.args, **self.kwargs))


__all__ = [
    "TimeOutError",
    "timeout",
    "timed_run",
    "context_timeout",
    "affirm",
    "ignored",
    "LazyEval",
]
